import cv2
import numpy as np
from image_processing import process_image_bytes, process_numpy_image
import time
import os

class ImageProcessor:
    def __init__(self):
        # 预加载模块，无需额外初始化
        pass
    
    def process(self, command: str, image: np.ndarray) -> np.ndarray:
        """处理图像并返回结果
        
        支持两种模式：
        1. 字节模式（PNG编码）
        2. Numpy数组模式（直接内存访问）
        """
        # 自动选择最优处理方式
        if image.size > 1024 * 1024:  # 大于1MB使用字节模式
            return self.process_bytes(command, image)
        else:
            return self.process_numpy(command, image)
    
    def process_bytes(self, command: str, image: np.ndarray) -> np.ndarray:
        """通过字节接口处理图像"""
        # 编码为PNG
        success, buffer = cv2.imencode(".png", image)
        if not success:
            raise ValueError("Failed to encode image as PNG")
        
        # 调用Rust处理
        start_time = time.perf_counter()
        result_bytes = process_image_bytes(buffer.tobytes(), command)
        elapsed = time.perf_counter() - start_time
        
        # 解码结果
        result_img = cv2.imdecode(
            np.frombuffer(result_bytes, dtype=np.uint8), 
            cv2.IMREAD_UNCHANGED
        )
        
        print(f"Processed '{command}' via bytes in {elapsed*1000:.2f}ms")
        return result_img
    
    def process_numpy(self, command: str, image: np.ndarray) -> np.ndarray:
        """通过numpy接口处理图像"""
        # 确保正确的通道顺序 (OpenCV使用BGR)
        if image.shape[2] == 3:
            image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        else:
            image_rgb = image
        
        # 调用Rust处理
        start_time = time.perf_counter()
        result_array = process_numpy_image(image_rgb, command)
        elapsed = time.perf_counter() - start_time
        
        # 转换回OpenCV格式
        if result_array.shape[2] == 3:
            result_img = cv2.cvtColor(result_array, cv2.COLOR_RGB2BGR)
        else:
            result_img = result_array
        
        print(f"Processed '{command}' via numpy in {elapsed*1000:.2f}ms")
        return result_img

def main():
    # 创建处理器
    processor = ImageProcessor()
    
    # 创建测试图像
    test_img = np.zeros((480, 640, 3), dtype=np.uint8)
    cv2.putText(test_img, "Test Image", (50, 240), 
                cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 255, 0), 3)
    
    # 测试所有处理命令
    commands = ["grayscale", "invert", "blur", "rotate90", "canny_edge"]
    
    for cmd in commands:
        try:
            # 处理图像
            result = processor.process(cmd, test_img)
            
            # 保存结果
            cv2.imwrite(f"result_{cmd}.jpg", result)
            print(f"Saved result for '{cmd}'")
        except Exception as e:
            print(f"Error processing '{cmd}': {str(e)}")
    
    # 性能测试
    print("\nPerformance test:")
    large_img = np.random.randint(0, 256, (1080, 1920, 3), dtype=np.uint8)
    
    for cmd in commands:
        start_time = time.perf_counter()
        for _ in range(10):
            _ = processor.process(cmd, large_img)
        elapsed = time.perf_counter() - start_time
        print(f"  {cmd}: {elapsed/10*1000:.2f}ms per image")
    
    print("All processing completed.")

if __name__ == "__main__":
    main()